import torch.nn as nn

from models.network.gcn_encoder import GCNEncoder
from models.text.tweet_encoder import TweetsEncoder


class TGAT(nn.Module):
    def __init__(self, config):
        super(TGAT, self).__init__()
        self.input_type = config['input_type']
        self.d_model = config['d_model']
        self.d_output = config['d_output']
        self.d_user_feat = config['d_user_feat']
        self.dataset = config['dataset']

        if self.input_type == 'text':
            self.tweets_encoder = TweetsEncoder(self.d_model)

        elif self.input_type == 'network':
            self.network_encoder = GCNEncoder(self.d_user_feat, d_model=self.d_model)

        elif self.input_type == 'all':
            self.tweets_encoder = TweetsEncoder(self.d_model)
            self.network_encoder = GCNEncoder(self.d_user_feat, d_model=self.d_model)
        else:
            raise Exception(
                "Parameters 'input_type' must be one of text, network or all, but get {}.".format(self.input_type))
        self.dropout = nn.Dropout(0.2)
        self.classifier = nn.Linear(self.d_model, self.d_output)
        self.cross_entropy = nn.CrossEntropyLoss()

    def fusion(self, text_encoding, network_encoding):
        return text_encoding + network_encoding

    def forward(self, graph, graph_feat, users, tweets, tweets_num):
        if self.input_type == 'text':
            text_encoding = self.tweets_encoder(tweets, tweets_num)
            encoding = text_encoding
        elif self.input_type == 'network':
            network_encoding = self.network_encoder(graph, graph_feat, users)
            encoding = network_encoding
        elif self.input_type == 'all':
            text_encoding = self.tweets_encoder(tweets, tweets_num)
            network_encoding = self.network_encoder(graph, graph_feat, users)
            encoding = self.fusion(text_encoding, network_encoding)

        encoding = self.dropout(encoding)
        output = self.classifier(encoding)
        return output

    def loss(self, preds, labels):
        return self.cross_entropy(preds, labels)
